package in.ac.iitb.cse.nlp.poswordpredictor;

import in.ac.iitb.cse.nlp.postagger.data.BiViterbiHMM;
import in.ac.iitb.cse.nlp.postagger.data.Config;
import in.ac.iitb.cse.nlp.postagger.data.DataMaps;
import in.ac.iitb.cse.nlp.postagger.data.EmissionMatrix;
import in.ac.iitb.cse.nlp.postagger.data.HMM;
import in.ac.iitb.cse.nlp.wordprediction.WordMatrix;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import jline.Completor;
import jline.ConsoleReader;

public class POSWordPredictor {
	ExtendedTransitionMatrix tm;
	EmissionMatrix em;
	InverseEmissionMatrix iem;
	WordMatrix wordMatrix;
	List<String> tagSet;
	Trainer trainer;
	ConsoleReader consoleReader;

	public POSWordPredictor() {
		trainer = newTraining();
	}

	public Trainer newTraining() {
		trainer = new Trainer();
		DataMaps.tags.clear();
		DataMaps.words.clear();
		trainer.goDir("", "corpus");
		tm = trainer.transitionMatrix;
		em = trainer.emissionMatrix;
		iem = trainer.inverseEmissionMatrix;
		tagSet = new ArrayList<String>(trainer.tagSet);
		Config.nTotalTags = trainer.tagSet.size();
		return trainer;
	}

	List<String> predict(String sentence) {
		List<Prediction> predictedWords = new ArrayList<Prediction>();
		List<Prediction> suitablePrediction = null;
		System.out.println();
		if (sentence.length() == 0) {
			suitablePrediction = tm
					.getSuitablePrediction(Constants.START_OF_SENTENCE);

		} else {
			List<String> tags = tagSentence(sentence);
			suitablePrediction = tm
					.getSuitablePrediction(tags.get(tags.size() - 1));
		}
		double factor = 1;
		for (Prediction prediction : suitablePrediction) {
			predictedWords.addAll(iem.getSuitablePrediction(prediction, factor));
			factor = factor - (1.0 / (double)Constants.MAX_NUM_OF_PREDICTONS);
			if (factor == 0)
				break;
		}
		String[] split = sentence.split(Trainer.WORD_SEPARATOR);
		predictedWords.addAll(trainer.wordMatrix
				.getSuitablePrediction(split[split.length - 1]));
		Collections.sort(predictedWords);
		List<String> finalPredictions = new ArrayList<String>();
		for (int i = 0; i < predictedWords.size()
				&& i < Constants.MAX_NUM_OF_PREDICTONS; i++) {
			finalPredictions.add(predictedWords.get(i).word);
		}
		return finalPredictions;
	}

	public List<String> tagSentence(String sentence) {

		String[] split = sentence.split(" ");
		List<String> input = new ArrayList<String>();
		for (String string : split) {
			if (string.isEmpty() == false)
				input.add(string);
		}

		HMM vhmm = null;
		vhmm = new BiViterbiHMM();
		vhmm.init(new ArrayList<String>(tagSet), tm, em);
		List<String> bestProbTagSeq = vhmm.bestProbTagSeq(input);
		if (bestProbTagSeq != null) {
			System.out.println("\n----------------------------------");
			for (int i = 0; i < bestProbTagSeq.size(); i++)
				System.out.print(bestProbTagSeq.get(i) + "_" + input.get(i)
						+ " ");
			System.out.println("\n----------------------------------");
		}
		return bestProbTagSeq;
	}

	public static void main(String[] args) {
		new POSWordPredictor().go();
	}

	private void go() {
		initialize();
	}

	private void initialize() {
		try {
			consoleReader = new ConsoleReader();
			consoleReader.addCompletor(new POSWordCompletor());
			String line = null;
			while ((line = consoleReader.readLine("pos_word_predictor>")) != null) {
				line=line.trim();
				if (line.equalsIgnoreCase("exit")
						|| line.equalsIgnoreCase("quit")) {
					System.exit(0);
				}
			}
		} catch (IOException e) {
			System.out.println("An error occured");
		}
	}

	class POSWordCompletor implements Completor {

		@Override
		public int complete(String buffer, int cursor, List completions) {
			String input = null;
			if (cursor == 0) {
				input = "";
			} else {
				input = buffer.trim();
			}
			for (String string : predict(input.trim())) {
				completions.add(" " + string);
			}
			return cursor;
		}

	}
}
